{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Code Generation and AutoDiff with a Cart-Pole System in Crocoddyl**\n",
    "\n",
    "In this notebook, we demonstrate how to perform code generation and compute automatic differentiation (AutoDiff) derivatives for an action model written in Python. Specifically, we develop an action model for a cart-pole system, similar to the one in the `cartpole_swingup.ipynb` notebook. The key features of this notebook are\n",
    " - **Understanding Code Generation**: Implement the cartpole dynamics and its analytical derivates using Crocoddyl's code generation framework.\n",
    " - **Automatic Differentiation**: Validate analytical derivatives with AutoDiff and compare computational performance.\n",
    " - **Optimal Control with Code-Generated Actions**: Solves an optimal control problem using the `SolveFDDP` solver and analyze the computational speedup from code generation.\n",
    "\n",
    "## Overview of the Cart-Pole System\n",
    "\n",
    "The cart-pole system consists of a cart that can move along a horizontal track and a pole attached to the cart via a pivot joint. The goal is to balance the pole in the upright position by applying forces to the cart.\n",
    "\n",
    "<div style=\"text-align:center;\">\n",
    "    <img src=\"https://www.ashwinnarayan.com/img/cart-pole.png\" width=\"300\" height=\"300\">\n",
    "</div>\n",
    "\n",
    "## Equations of Motion\n",
    "\n",
    "The state of the system is given by:\n",
    "\n",
    "- $y$: Position of the cart\n",
    "- $\\theta$: Angle of the pole from the upright position\n",
    "- $\\dot{y}$: Velocity of the cart\n",
    "- $\\dot{\\theta}$: Angular velocity of the pole\n",
    "\n",
    "For more details on the equation of motions, integrator, and cost functions please go to the `cartpole_swingup.ipynb` notebook.\n",
    "\n",
    "## Code Generation in Crocoddyl\n",
    "\n",
    "When enabling `BUILD_WITH_CODEGEN_SUPPORT`, Crocoddyl leverages `CppADCodeGen` for efficient code generation and automatic differentiation of action models. To enable this, we modify the `ActionModelCartpole` from `cartpole_swingup.ipynb` as follows:\n",
    " 1. Inheritance: Extend Crocoddyl’s base action model supporting `CppAD<CppAD::cg::CG<>>` types via the `cgfloat64` module.\n",
    " 2. Function Replacement: Replace unsupported primitive functions with `CppADCodeGen` equivalents (e.g., `numpy.sin` → `pycppad.sin_me`).\n",
    "\n",
    "The following section presents the modified cart-pole action model in detail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import crocoddyl.cgfloat64\n",
    "from pycppad import ADCG, CG\n",
    "import numpy as np\n",
    "\n",
    "class ActionModelCartpole(crocoddyl.cgfloat64.ActionModelAbstract):\n",
    "    \"\"\"\n",
    "    Defines the action model for the cart-pole system in Crocoddyl.\n",
    "    \n",
    "    This model simulates a cart moving along a horizontal track with a \n",
    "    freely rotating pole attached to it. The goal is to control the force \n",
    "    applied to the cart to stabilize the pole.\n",
    "    \n",
    "    State vector: [y, θ, ẏ, θ̇]\n",
    "    Control input: [f] (force applied to the cart)\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        Initializes the cart-pole action model with system parameters and cost weights.\n",
    "        \"\"\"\n",
    "        crocoddyl.cgfloat64.ActionModelAbstract.__init__(\n",
    "            self, crocoddyl.cgfloat64.StateVector(4), nu=1, nr=6)\n",
    "\n",
    "        self.Δt = ADCG(CG(5e-2))\n",
    "        self.m_cart = ADCG(CG(1.0))\n",
    "        self.m_pole = ADCG(CG(0.1))\n",
    "        self.l_pole = ADCG(CG(0.5))\n",
    "        self.grav = ADCG(CG(9.81))\n",
    "        self.costWeights = [\n",
    "            ADCG(CG(1.0)),\n",
    "            ADCG(CG(1.0)),\n",
    "            ADCG(CG(0.1)),\n",
    "            ADCG(CG(0.001)),\n",
    "            ADCG(CG(0.001)),\n",
    "            ADCG(CG(1.0)),\n",
    "        ]  # sin, 1-cos, x, xdot, thdot, f\n",
    "\n",
    "    def calc(self, data, x, u=None):\n",
    "        \"\"\"\n",
    "        Computes the next state and the cost function.\n",
    "\n",
    "        Args:\n",
    "            data (ActionDataCartpole): Data structure to store intermediate results.\n",
    "            x (numpy.ndarray): State vector [y, θ, ẏ, θ̇].\n",
    "            u (numpy.ndarray, optional): Control input [f]. Defaults to None.\n",
    "\n",
    "        Updates:\n",
    "            - Next state xnext.\n",
    "            - Cost residuals r.\n",
    "            - Total cost value.\n",
    "        \"\"\"\n",
    "        if u is not None:\n",
    "            # Getting the state and control variables\n",
    "            y, θ, ẏ, θ̇, f = x[0], x[1], x[2], x[3], u[0]\n",
    "            # Shortname for system parameters\n",
    "            Δt, m_cart, m_pole, l_pole, grav, w = self.Δt, self.m_cart, self.m_pole, self.l_pole, self.grav, self.costWeights\n",
    "            sin_θ, cos_θ = ADCG.sin_me(θ), ADCG.cos_me(θ)\n",
    "            # Computing the cartpole dynamics\n",
    "            data.μ = m_cart + m_pole * sin_θ* sin_θ\n",
    "            data.ÿ̈ = f / data.μ\n",
    "            data.ÿ̈ += m_pole * grav * cos_θ * sin_θ / data.μ\n",
    "            data.ÿ̈ -= m_pole * l_pole * (sin_θ * θ̇ * θ̇ / data.μ)\n",
    "            data.θ̈ = (f / l_pole) * cos_θ / data.μ\n",
    "            data.θ̈ += ((m_cart + m_pole) * grav / l_pole) * sin_θ / data.μ\n",
    "            data.θ̈ -= m_pole * cos_θ * sin_θ * θ̇ * θ̇ / data.μ\n",
    "            data.ẏ_next = ẏ + Δt * data.ÿ̈\n",
    "            data.θ̇_next = θ̇ + Δt * data.θ̈\n",
    "            data.y_next = y + Δt * data.ẏ_next\n",
    "            data.θ_next = θ + Δt * data.θ̇_next\n",
    "            data.xnext[:] = np.array([data.y_next, data.θ_next, data.ẏ_next, data.θ̇_next])\n",
    "            # Computing the cost residual and value\n",
    "            data.r[:] = w * np.array([sin_θ, ADCG(CG(1.0)) - cos_θ, y, ẏ, θ̇, f])\n",
    "            data.cost = ADCG(CG(0.0))\n",
    "            for i in range(self.nr):\n",
    "                data.cost += ADCG(CG(0.5)) * (data.r[i] * data.r[i])\n",
    "        else:\n",
    "            # Getting the state and control variables\n",
    "            y, θ, ẏ, θ̇ = x[0], x[1], x[2], x[3]\n",
    "            w = self.costWeights\n",
    "            sin_θ, cos_θ = ADCG.sin_me(θ), ADCG.cos_me(θ)\n",
    "            data.xnext[:] = x\n",
    "            # Computing the cost residual and value\n",
    "            data.r[:] = w * np.array([sin_θ, ADCG(CG(1.0)) - cos_θ, y, ẏ, θ̇, ADCG(CG(0.0))])\n",
    "            data.cost = ADCG(CG(0.0))\n",
    "            for i in range(self.nr):\n",
    "                data.cost += ADCG(CG(0.5)) * (data.r[i] * data.r[i])\n",
    "\n",
    "    def calcDiff(self, data, x, u=None):\n",
    "        \"\"\"\n",
    "        Computes the derivatives of the dynamics and the cost function.\n",
    "\n",
    "        Args:\n",
    "            data (ActionDataCartpole): Data structure to store intermediate results.\n",
    "            x (numpy.ndarray): State vector [y, θ, ẏ, θ̇].\n",
    "            u (numpy.ndarray, optional): Control input [f]. Defaults to None.\n",
    "\n",
    "        Updates:\n",
    "            - Derivetives of the dynamics.\n",
    "            - Derivatives of the cost function.\n",
    "        \"\"\"\n",
    "        if u is not None:\n",
    "            # Getting the state and control variables\n",
    "            y, θ, ẏ, θ̇, f = x[0], x[1], x[2], x[3], u[0]\n",
    "            # Shortname for system parameters\n",
    "            Δt, m_cart, m_pole, l_pole, grav, w = self.Δt, self.m_cart, self.m_pole, self.l_pole, self.grav, self.costWeights\n",
    "            sin_θ, cos_θ = ADCG.sin_me(θ), ADCG.cos_me(θ)\n",
    "            # Computing the derivative of the cartpole dynamics\n",
    "            data.dμ_dθ = ADCG(CG(2.0)) * m_pole * sin_θ * cos_θ\n",
    "            data.dÿ̈_dy = ADCG(CG(0.0))\n",
    "            data.dÿ̈_dθ = m_pole * grav * (cos_θ * cos_θ - sin_θ * sin_θ) / data.μ\n",
    "            data.dÿ̈_dθ -= m_pole * l_pole * cos_θ * θ̇ * θ̇ / data.μ\n",
    "            data.dÿ̈_dθ -= data.dμ_dθ * data.ÿ̈ / data.μ\n",
    "            data.dÿ̈_dẏ = ADCG(CG(0.0))\n",
    "            data.dÿ̈_dθ̇ = -ADCG(CG(2.0)) * m_pole * l_pole * sin_θ * θ̇ / data.μ\n",
    "            data.dÿ̈_du = ADCG(CG(1.0)) / data.μ\n",
    "            data.dθ̈_dy = ADCG(CG(0.0))\n",
    "            data.dθ̈_dθ = -(f / l_pole) * sin_θ / data.μ\n",
    "            data.dθ̈_dθ += ((m_cart + m_pole) * grav / l_pole) * cos_θ / data.μ\n",
    "            data.dθ̈_dθ -= m_pole * (cos_θ * cos_θ - sin_θ * sin_θ) * θ̇ * θ̇ / data.μ\n",
    "            data.dθ̈_dθ -= data.dμ_dθ * data.θ̈ / data.μ\n",
    "            data.dθ̈_dẏ = ADCG(CG(0.0))\n",
    "            data.dθ̈_dθ̇ = -ADCG(CG(2.0)) * m_pole * cos_θ * sin_θ * θ̇ / data.μ\n",
    "            data.dθ̈_du = cos_θ / (l_pole * data.μ)\n",
    "            data.dẏ_next_dy = Δt * data.dÿ̈_dy\n",
    "            data.dẏ_next_dθ = Δt * data.dÿ̈_dθ\n",
    "            data.dẏ_next_dẏ = ADCG(CG(1.0)) + Δt * data.dÿ̈_dẏ\n",
    "            data.dẏ_next_dθ̇ = Δt * data.dÿ̈_dθ̇\n",
    "            data.dẏ_next_du = Δt * data.dÿ̈_du\n",
    "            data.dθ̇_next_dy = Δt * data.dθ̈_dy\n",
    "            data.dθ̇_next_dθ = Δt * data.dθ̈_dθ\n",
    "            data.dθ̇_next_dẏ = Δt * data.dθ̈_dẏ\n",
    "            data.dθ̇_next_dθ̇ = ADCG(CG(1.0)) + Δt * data.dθ̈_dθ̇\n",
    "            data.dθ̇_next_du = Δt * data.dθ̈_du\n",
    "            data.dy_next_dy = ADCG(CG(1.0)) + Δt * data.dẏ_next_dy\n",
    "            data.dy_next_dθ = Δt * data.dẏ_next_dθ\n",
    "            data.dy_next_dẏ = Δt * data.dẏ_next_dẏ\n",
    "            data.dy_next_dθ̇ = Δt * data.dẏ_next_dθ̇\n",
    "            data.dy_next_du = Δt * data.dẏ_next_du\n",
    "            data.dθ_next_dy = Δt * data.dθ̇_next_dy\n",
    "            data.dθ_next_dθ = ADCG(CG(1.0)) + Δt * data.dθ̇_next_dθ\n",
    "            data.dθ_next_dẏ = Δt * data.dθ̇_next_dẏ\n",
    "            data.dθ_next_dθ̇ = Δt * data.dθ̇_next_dθ̇\n",
    "            data.dθ_next_du = Δt * data.dθ̇_next_du\n",
    "            data.Fx[:, :] = np.array([[data.dy_next_dy, data.dy_next_dθ, data.dy_next_dẏ, data.dy_next_dθ̇],\n",
    "                                      [data.dθ_next_dy, data.dθ_next_dθ, data.dθ_next_dẏ, data.dθ_next_dθ̇],\n",
    "                                      [data.dẏ_next_dy, data.dẏ_next_dθ, data.dẏ_next_dẏ, data.dẏ_next_dθ̇],\n",
    "                                      [data.dθ̇_next_dy, data.dθ̇_next_dθ, data.dθ̇_next_dẏ, data.dθ̇_next_dθ̇]])\n",
    "            data.Fu[:] = np.array([data.dy_next_du, data.dθ_next_du, data.dẏ_next_du, data.dθ̇_next_du])\n",
    "            # Computing derivatives of the cost function\n",
    "            w0_2, w1_2, w2_2, w3_2, w4_2, w5_2 = w[0] * w[0], w[1] * w[1], w[2] * w[2], w[3] * w[3], w[4] * w[4], w[5] * w[5]\n",
    "            data.Lx[0] = w2_2 * y\n",
    "            data.Lx[1] = w0_2 * sin_θ * cos_θ + w1_2 * (ADCG(CG(1.0)) - cos_θ) * sin_θ\n",
    "            data.Lx[2] = w3_2 * ẏ\n",
    "            data.Lx[3] = w4_2 * θ̇\n",
    "            data.Lu[0] = w5_2 * f\n",
    "            data.Lxx[0, 0] = w2_2\n",
    "            data.Lxx[1, 1] = w0_2 * (cos_θ * cos_θ - sin_θ * sin_θ)\n",
    "            data.Lxx[1, 1] += w1_2 * ((ADCG(CG(1.0)) - cos_θ) * cos_θ + sin_θ * sin_θ)\n",
    "            data.Lxx[2, 2] = w3_2\n",
    "            data.Lxx[3, 3] = w4_2\n",
    "            data.Luu[:] = w5_2\n",
    "        else:\n",
    "            # Getting the state and control variables\n",
    "            y, θ, ẏ, θ̇ = x[0], x[1], x[2], x[3]\n",
    "            w = self.costWeights\n",
    "            sin_θ, cos_θ = ADCG.sin_me(θ), ADCG.cos_me(θ)\n",
    "            # Computing the derivative of the cartpole dynamics\n",
    "            for i in range(self.state.ndx):\n",
    "                data.Fx[i, i] = ADCG(CG(1.0))\n",
    "            # Computing derivatives of the cost function\n",
    "            w0_2, w1_2, w2_2, w3_2, w4_2, w5_2 = w[0] * w[0], w[1] * w[1], w[2] * w[2], w[3] * w[3], w[4] * w[4], w[5] * w[5]\n",
    "            data.Lx[0] = w2_2 * y\n",
    "            data.Lx[1] = w0_2 * sin_θ * cos_θ + w1_2 * (ADCG(CG(1.0)) - cos_θ) * sin_θ\n",
    "            data.Lx[2] = w3_2 * ẏ\n",
    "            data.Lx[3] = w4_2 * θ̇\n",
    "            data.Lxx[0, 0] = w2_2\n",
    "            data.Lxx[1, 1] = w0_2 * (cos_θ * cos_θ - sin_θ * sin_θ)\n",
    "            data.Lxx[1, 1] += w1_2 * ((ADCG(CG(1.0)) - cos_θ) * cos_θ + sin_θ * sin_θ)\n",
    "            data.Lxx[2, 2] = w3_2\n",
    "            data.Lxx[3, 3] = w4_2\n",
    "\n",
    "    def createData(self):\n",
    "        \"\"\"\n",
    "        Creates the action data structure for the cart-pole model.\n",
    "\n",
    "        Returns:\n",
    "            ActionDataCartpole: Data structure to store intermediate computations.\n",
    "        \"\"\"\n",
    "        return ActionDataCartpole(self)\n",
    "\n",
    "\n",
    "class ActionDataCartpole(crocoddyl.cgfloat64.ActionDataAbstract):\n",
    "    \"\"\"\n",
    "    Data structure for storing intermediate computations of the cart-pole dynamics.\n",
    "    \"\"\"\n",
    "    def __init__(self, model):\n",
    "        \"\"\"\n",
    "        Initializes the data structure with default values.\n",
    "\n",
    "        Args:\n",
    "            model (ActionModelCartpole): Cart-pole action model.\n",
    "        \"\"\"\n",
    "        crocoddyl.cgfloat64.ActionDataAbstract.__init__(self, model)\n",
    "        self.μ = ADCG(CG(0.0))\n",
    "        self.ÿ̈ = ADCG(CG(0.0))\n",
    "        self.θ̈ = ADCG(CG(0.0))\n",
    "        self.ẏ_next = ADCG(CG(0.0))\n",
    "        self.θ̇_next = ADCG(CG(0.0))\n",
    "        self.y_next = ADCG(CG(0.0))\n",
    "        self.θ_next = ADCG(CG(0.0))\n",
    "        self.dμ_dθ = ADCG(CG(0.0))\n",
    "        self.dÿ̈_dy = ADCG(CG(0.0))\n",
    "        self.dÿ̈_dθ = ADCG(CG(0.0))\n",
    "        self.dÿ̈_dθ = ADCG(CG(0.0))\n",
    "        self.dÿ̈_dθ = ADCG(CG(0.0))\n",
    "        self.dÿ̈_dẏ = ADCG(CG(0.0))\n",
    "        self.dÿ̈_dθ̇ = ADCG(CG(0.0))\n",
    "        self.dÿ̈_du = ADCG(CG(0.0))\n",
    "        self.dθ̈_dy = ADCG(CG(0.0))\n",
    "        self.dθ̈_dθ = ADCG(CG(0.0))\n",
    "        self.dθ̈_dθ = ADCG(CG(0.0))\n",
    "        self.dθ̈_dθ = ADCG(CG(0.0))\n",
    "        self.dθ̈_dθ = ADCG(CG(0.0))\n",
    "        self.dθ̈_dẏ = ADCG(CG(0.0))\n",
    "        self.dθ̈_dθ̇ = ADCG(CG(0.0))\n",
    "        self.dẏ_next_dy = ADCG(CG(0.0))\n",
    "        self.dẏ_next_dθ = ADCG(CG(0.0))\n",
    "        self.dẏ_next_dẏ = ADCG(CG(0.0))\n",
    "        self.dẏ_next_dθ̇ = ADCG(CG(0.0))\n",
    "        self.dẏ_next_du = ADCG(CG(0.0))\n",
    "        self.dθ̇_next_dy = ADCG(CG(0.0))\n",
    "        self.dθ̇_next_dθ = ADCG(CG(0.0))\n",
    "        self.dθ̇_next_dẏ = ADCG(CG(0.0))\n",
    "        self.dθ̇_next_dθ̇ = ADCG(CG(0.0))\n",
    "        self.dθ̇_next_du = ADCG(CG(0.0))\n",
    "        self.dy_next_dy = ADCG(CG(0.0))\n",
    "        self.dy_next_dθ = ADCG(CG(0.0))\n",
    "        self.dy_next_dẏ = ADCG(CG(0.0))\n",
    "        self.dy_next_dθ̇ = ADCG(CG(0.0))\n",
    "        self.dy_next_du = ADCG(CG(0.0))\n",
    "        self.dθ_next_dy = ADCG(CG(0.0))\n",
    "        self.dθ_next_dθ = ADCG(CG(0.0))\n",
    "        self.dθ_next_dẏ = ADCG(CG(0.0))\n",
    "        self.dθ_next_dθ̇ = ADCG(CG(0.0))\n",
    "        self.dθ_next_du = ADCG(CG(0.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validating Analytical Derivatives with Automatic Differentiation (AutoDiff)\n",
    "\n",
    "Once the action model is defined and its analytical derivatives for the dynamics and cost function are derived, it is crucial to verify their correctness. Since our action model supports code generation, we validate these derivatives using **automatic differentiation (AutoDiff)**. AutoDiff computes the Jacobian and Hessian by directly evaluating the expression and then generates an optimized C++ library, which is linked within `ActionModelCodeGen`.\n",
    "\n",
    "Generating a C++ library through code generation can take a few seconds, depending on the complexity of the expressions. This process can be performed with or without AutoDiff. Below, we demonstrate how to generate action models both ways."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code generate ActionModelCarpole without and with AutoDiff\n",
    "model = crocoddyl.ActionModelCodeGen(ActionModelCartpole(), \"cartpole\")\n",
    "data = model.createData()\n",
    "model_ad = crocoddyl.ActionModelCodeGen(ActionModelCartpole(), \"cartpole_ad\", True)\n",
    "data_ad = model_ad.createData()\n",
    "\n",
    "# Generate a random state and control\n",
    "x = model.state.rand()\n",
    "u = np.random.rand(1)\n",
    "\n",
    "# Compute analytical and NumDiff derivatives\n",
    "model.calc(data, x, u)\n",
    "model.calcDiff(data, x, u)\n",
    "model_ad.calc(data_ad, x, u)\n",
    "model_ad.calcDiff(data_ad, x, u)\n",
    "\n",
    "# Compare derivatives\n",
    "print(\"Fx difference:\", np.linalg.norm(data.Fx - data_ad.Fx))\n",
    "print(\"Fu difference:\", np.linalg.norm(data.Fu - data_ad.Fu))\n",
    "print(\"Lx difference:\", np.linalg.norm(data.Lx - data_ad.Lx))\n",
    "print(\"Lu difference:\", np.linalg.norm(data.Lu - data_ad.Lu))\n",
    "print(\"Lxx difference:\", np.linalg.norm(data.Lxx - data_ad.Lxx))\n",
    "print(\"Lxu difference:\", np.linalg.norm(data.Lxu - data_ad.Lxu))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking Code-Generated Analytical vs. Automatic Differentiation\n",
    "\n",
    "To assess the computational cost of AutoDiff, we compare its performance against code-generated analytical derivatives. Specifically, we run `calc` and `calcDiff` 1 millon times and profile the execution using Crocoddyl’s built-in profiling framework. The process is straightforward, as demonstrated below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "# Number of trials\n",
    "n_iter = 1000000\n",
    "\n",
    "# Enable and reset the profiler\n",
    "crocoddyl.enable_profiler()\n",
    "crocoddyl.stop_watch_reset_all()\n",
    "\n",
    "# Run calc and calcDiff without AutoDiff\n",
    "for _ in range(n_iter):\n",
    "    model.calc(data, x, u)\n",
    "    model.calcDiff(data, x, u)\n",
    "\n",
    "# Report the profiling results for CodeGen\n",
    "crocoddyl.stop_watch_report(4)\n",
    "sys.stdout.flush()\n",
    "\n",
    "# Store the computation time associated to calc and calcDiff\n",
    "t_calc = crocoddyl.stop_watch_get_total_time(\"ActionModelCodeGen::calc\")\n",
    "t_calcDiff = crocoddyl.stop_watch_get_total_time(\"ActionModelCodeGen::calcDiff\")\n",
    "\n",
    "# Enable and reset the profiler\n",
    "crocoddyl.stop_watch_reset_all()\n",
    "crocoddyl.enable_profiler()\n",
    "\n",
    "# Run calc and calcDiff without AutoDiff\n",
    "for _ in range(n_iter):\n",
    "    model_ad.calc(data_ad, x, u)\n",
    "    model_ad.calcDiff(data_ad, x, u)\n",
    "\n",
    "# Store the computation time associated to calc and calcDiff\n",
    "t_calc_ad = crocoddyl.stop_watch_get_total_time(\"ActionModelCodeGen::calc\")\n",
    "t_calcDiff_ad = crocoddyl.stop_watch_get_total_time(\"ActionModelCodeGen::calcDiff\")\n",
    "\n",
    "# Report the profiling results for AutoDiff\n",
    "crocoddyl.stop_watch_report(4)\n",
    "sys.stdout.flush()\n",
    "\n",
    "ration_cg = t_calc_ad / t_calc\n",
    "ration = t_calcDiff_ad / t_calcDiff\n",
    "print()\n",
    "print(f\"Code generation with analytical derivatives is {ration_cg:.2f} times as fast as AutoDiff.\")\n",
    "print(f\"Analytical derivatives are {ration:.2f} times as fast as AutoDiff.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Balancing our Cart-Pole System using Optimal Control\n",
    "\n",
    "After verifying the analytical derivatives, we now solve the optimal control problem introduced in `cartpole_swingup.ipynb`, but this time using a code-generated library for computing cost, dynamics, and their derivatives.\n",
    "\n",
    "It is worth noting that Crocoddyl’s Code Generation and AutoDiff framework supports both equality and inequality constraints. However, for simplicity, we ignore them in this example.\n",
    "\n",
    "The optimal control problem is formulated as follows:\n",
    "\n",
    "\\begin{aligned}\n",
    "& \\underset{u(t)}{\\text{minimize}} \\quad J = \\frac{1}{2} \\sum_{i=0}^{N} w_i \\, r_i^2(\\mathbf{x}, \\mathbf{u})  \\\\\n",
    "& \\text{subject to} \\quad \\mathbf{x}^+ = f(\\mathbf{x}, \\mathbf{u})\n",
    "\\end{aligned}\n",
    "\n",
    "Where:\n",
    " - $J$ is the cost function, which sums over the trajectory nodes  i  the weighted square of the residual  r_i(\\mathbf{x}, \\mathbf{u}) , capturing the deviation from the desired state and control input.\n",
    " - $w_i$ is the weight associated with each node in the trajectory, typically representing how much importance we assign to a given time step in the optimization.\n",
    " - $\\mathbf{x}^+$  represents the state at the next time step, while $\\mathbf{x}$ is the current state.\n",
    " - $f(\\mathbf{x}, \\mathbf{u})$ represents the system’s dynamics, describing how the state evolves from one time step to the next under the control input $\\mathbf{u}[i]$.\n",
    "\n",
    "The implementation is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "from cartpole_utils import animateCartpole\n",
    "import sys\n",
    "\n",
    "# Create a shooting problem with 50 running nodes (i.e., 50 * 5e-2 = 2.5 sec)\n",
    "N = 50\n",
    "x0 = np.array([0.0, 0.0, 1.0, 0.5])\n",
    "problem = crocoddyl.ShootingProblem(x0, [model.copy() for _ in range(N)], model)\n",
    "\n",
    "# Creating the FDDP solver and setting the logger callback\n",
    "solver = crocoddyl.SolverFDDP(problem)\n",
    "solver.setCallbacks([crocoddyl.CallbackVerbose()])\n",
    "\n",
    "crocoddyl.enable_profiler()\n",
    "\n",
    "# Solving this problem and display the optimal trajectory of our cart-pole system\n",
    "solver.solve()\n",
    "\n",
    "crocoddyl.stop_watch_report(4)\n",
    "crocoddyl.stop_watch_reset_all()\n",
    "sys.stdout.flush()\n",
    "\n",
    "anim = animateCartpole(solver.xs)\n",
    "# HTML(anim.to_jshtml())\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Changing the Reference in Our Cart-Pole Action Model\n",
    "\n",
    "In many applications, we need to dynamically update reference values. A common example is **model predictive control (MPC)**, where cost references or constraint activations change over time.\n",
    "\n",
    "For a **code-generated action model**, we handle this by defining a *parameter environment*. This environment determines how reference vectors modify internal parameters within the action model. It is implemented through a function that updates the relevant parameters accordingly.\n",
    "\n",
    "In our cart-pole example, we aim to adjust the cost weights dynamically, as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the parameters environemnt\n",
    "def cost_weights_env(model, vector):\n",
    "    model.costWeights[:] = vector\n",
    "\n",
    "# Code generate ActionModelCarpole with parameters\n",
    "model_env = crocoddyl.ActionModelCodeGen(ActionModelCartpole(), \"cartpole\", False, 6, cost_weights_env)\n",
    "data_env = model_env.createData()\n",
    "\n",
    "# Define the same parameters as before\n",
    "model_env.update_p(data_env, np.array([1.0, 1.0, 0.1, 0.001, 0.001, 1.0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Introducing parameters in our code-generated action model typically has **minimal impact on computation time**, as demonstrated in the following benchmark:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "# Number of trials\n",
    "n_iter = 1000000\n",
    "\n",
    "# Enable and reset the profiler\n",
    "crocoddyl.enable_profiler()\n",
    "crocoddyl.stop_watch_reset_all()\n",
    "\n",
    "# Run calc and calcDiff without parameters\n",
    "for _ in range(n_iter):\n",
    "    model.calc(data, x, u)\n",
    "    model.calcDiff(data, x, u)\n",
    "\n",
    "# Report the profiling results for parameters\n",
    "crocoddyl.stop_watch_report(4)\n",
    "sys.stdout.flush()\n",
    "\n",
    "# Store the computation time associated to calc and calcDiff\n",
    "t_calc = crocoddyl.stop_watch_get_total_time(\"ActionModelCodeGen::calc\")\n",
    "t_calcDiff = crocoddyl.stop_watch_get_total_time(\"ActionModelCodeGen::calcDiff\")\n",
    "\n",
    "# Enable and reset the profiler\n",
    "crocoddyl.stop_watch_reset_all()\n",
    "crocoddyl.enable_profiler()\n",
    "\n",
    "# Run calc and calcDiff without AutoDiff\n",
    "for _ in range(n_iter):\n",
    "    model_env.calc(data_env, x, u)\n",
    "    model_env.calcDiff(data_env, x, u)\n",
    "\n",
    "# Store the computation time associated to calc and calcDiff\n",
    "t_calc_ad = crocoddyl.stop_watch_get_total_time(\"ActionModelCodeGen::calc\")\n",
    "t_calcDiff_ad = crocoddyl.stop_watch_get_total_time(\"ActionModelCodeGen::calcDiff\")\n",
    "\n",
    "# Report the profiling results for AutoDiff\n",
    "crocoddyl.stop_watch_report(4)\n",
    "sys.stdout.flush()\n",
    "\n",
    "ration_cg = t_calc / t_calc_ad * 100\n",
    "ration = t_calcDiff / t_calcDiff_ad * 100\n",
    "print()\n",
    "print(f\"Code generation with analytical derivatives is {ration_cg:.2f}% times faster than AutoDiff.\")\n",
    "print(f\"Analytical derivatives are {ration:.2f}% times faster than AutoDiff.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we update the cost weights for the terminal node. As expected, this modification alters the behavior compared to the previous solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "from cartpole_utils import animateCartpole\n",
    "import sys\n",
    "\n",
    "# Create a shooting problem with 50 running nodes (i.e., 50 * 5e-2 = 2.5 sec)\n",
    "N = 50\n",
    "x0 = np.array([0.0, 0.0, 1.0, 0.5])\n",
    "problem = crocoddyl.ShootingProblem(x0, [model_env.copy() for _ in range(N)], model_env)\n",
    "\n",
    "# Change the cost weight associated to the terminal node\n",
    "problem.terminalModel.update_p(problem.terminalData, np.array([100, 100., 1., 0.1, 0.01, 0.001]))\n",
    "\n",
    "# Creating the FDDP solver and setting the logger callback\n",
    "solver = crocoddyl.SolverFDDP(problem)\n",
    "solver.setCallbacks([crocoddyl.CallbackVerbose()])\n",
    "\n",
    "crocoddyl.enable_profiler()\n",
    "\n",
    "# Solving this problem and display the optimal trajectory of our cart-pole system\n",
    "solver.solve()\n",
    "\n",
    "crocoddyl.stop_watch_report(4)\n",
    "crocoddyl.stop_watch_reset_all()\n",
    "sys.stdout.flush()\n",
    "\n",
    "anim = animateCartpole(solver.xs)\n",
    "# HTML(anim.to_jshtml())\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What’s Next?\n",
    "\n",
    "Now that we’ve explored how to solve an optimal control problem using FDDP and analyzed the benefits of code generation, the next step is to extend these concepts to more **general robotic systems**. In this notebook, we will cover:\n",
    "\n",
    "1. **Action Model Code Generation**: Learn how to automatically code generate of complex action models needed to model robotics via the `ActionModelCodeGen` class.\n",
    "\n",
    "2. **Derivatives and Sensitivity Analysis**: Validate code-generated derivatives by comparing them against code-generated automatic derivatives (e.g., using AutoDiff).\n",
    "\n",
    "3. **Profiling and Performance Evaluation**: Compare the performance gains from AutoDiff derivatives vs. manually defined derivatives using code generation.\n",
    "\n",
    "4. **Scaling to More Complex Systems**:\tUnderstand how adapt reference in a code-generated action model.\n",
    "\n",
    "By the end of the next notebook `scaling_to_robotics.ipynb`, you’ll have the skills to generate, validate, and optimize action models for your own robotics applications—unlocking faster and more accurate motion planning and control. 🚀"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "crocoddyl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
